from snort_rules import *
import math
import random
import string
import exrex
path = "/root/github/internet_product_safe_test/analyze_of_complexity"
#alert tcp $EXTERNAL_NET any -> $SMTP_SERVERS 25 ( msg:"rule_attrROWSER-CHROME Google Chrome NotifyInstanceWasDeleted
# emrule_attred use after free attempt"; flow:to_server,estarule_attrlished;
# file_data; content:"<emrule_attred ",nocase; content:"application/pdf",within 125,nocase;
# content:".addEventListener",nocase; content:".innerHTML",nocase;
# content:".reload(",nocase; content:"fromCharCode",fast_pattern,nocase;
# pcre:"/\.innerHTML\s*?\x3d\s*?[\x22\x27]\s*?[\x22\x27]/smi";
# metadata:policy max-detect-ips drop; service:smtp; reference:rule_attrugtraq,62752; reference:cve,2013-2912; classtype:attempted-user; sid:31599; rev:4; )

# a=open("./expRules.rules")
# for line in a:
#     print(line)

class ConstructHardMatchData(object):
    """
    Take in one Snort rule[class SnortRuleAttr] and a looptime number[Default 150] construct a hard match data for fuzzing
    """

    def __init__(self,snortrule,looptime=300):
        self.rule = snortrule
        self.snortrule_attr=SnortRuleAttr(self.rule)
        self.content_and_pcre=[]
        self.converge_content_and_pcre()
        self.data=""
        self.datalen=0
        self.looptime=looptime

    def converge_content_and_pcre(self):
        self.content_and_pcre = self.snortrule_attr.get_opt_content()
        pcre_index = self.snortrule_attr.option_index('pcre')[0]
        index = -1
        add_flag = False
        for num in self.snortrule_attr.option_index('content'):
            index = index + 1
            if num < pcre_index:
                continue
            else:
                self.content_and_pcre.insert(index, self.snortrule_attr.get_opt_pcre())
                add_flag = True
        if not add_flag:
            self.content_and_pcre.append(self.snortrule_attr.get_opt_pcre())
        # for case in self.content_and_pcre:
        #     print(case)
    def show_content_and_pcre(self):
        for case in self.content_and_pcre:
            print(case)

    def construct_data(self):
        length=len(self.content_and_pcre)
        index=0
        while index <= length-2:
            #basic information gather
            current_dic=self.content_and_pcre[index]
            next_dic=self.content_and_pcre[index+1]
            next_has_within = False
            next_has_distance = False
            current_has_depth=False
            current_has_offset=False
            if 'within' in next_dic and next_dic['within'] !=None:
                next_has_within=True   #current increase should less than within
            if 'distance' in next_dic and next_dic['distance'] !=None:
                next_has_distance=True  #for later use
            if 'depth' in next_dic and current_dic['depth'] !=None:
                current_has_depth=True  #for later use
            if 'offset' in next_dic and current_dic['offset'] !=None:
                current_has_offset=True  #for later use

            if 'match' in current_dic:
                mode='content'
            else:mode='pcre'
            ceil_loop_time=self.looptime   #[TODO]  Default 150
            floor_loop_time=1
            #caculate proper current_dic loop time
            if mode=='content':
                seed=current_dic['match']
                seed_len=len(seed)
                if next_has_within:
                    max_loop_time=self.max_loop_time_decide_by_within(seed_len,next_dic['within'])
                    if max_loop_time<ceil_loop_time:
                        ceil_loop_time=max_loop_time
                if next_has_distance:
                    min_loop_time=self.min_loop_time_decide_by_distance(seed_len,next_dic['distance'])
                    if min_loop_time>floor_loop_time:
                        floor_loop_time=min_loop_time
                if current_has_depth:
                    max_loop_time=self.max_loop_time_decide_by_depth(self.datalen,seed_len,current_dic['depth'])
                    if max_loop_time==None: pass
                    if max_loop_time<ceil_loop_time:
                        ceil_loop_time=max_loop_time
                if current_has_offset:
                    status=self.min_loop_time_decide_by_offset(self.datalen,seed_len,current_dic['offset'])
                    if status==False:   #means need to pack the blank
                        self.data=self.data+''.join(random.sample(string.ascii_letters + string.digits, self.datalen-current_dic['offset']))
                #loop_time=random.randint(ceil_loop_time,ceil_loop_time) #[TODO] both are ceil_loop_time
                final_loop_time=ceil_loop_time
                self.data=self.data+self.loop_seed(seed,final_loop_time)
                self.datalen=len(self.data)
            if mode=='pcre':
                pass  # presume pcre is always the last one in content_and_pcre
            index=index+1
        # self.data = self.data + ''.join(
        #     random.sample(string.ascii_letters + string.digits, 10))

    def data_add_last_pcre(self):
        """
        Add the pcre part to make the data  match the rule
        """
        pcre=self.content_and_pcre[len(self.content_and_pcre)-1]['pattern']
        self.data=self.data+exrex.getone(pcre)


    def max_loop_time_decide_by_within(self,seed_len,next_within_number):
        """
        Decide current seed max loop time by current seed length
        and next seed 'within' number
        """
        return math.floor(next_within_number/seed_len)+1

    def min_loop_time_decide_by_distance(self,seed_len,next_distance_number):
        """
        Decide current seed min loop time by current seed length
        and next seed 'distance' number
        """
        return math.ceil(next_distance_number/seed_len)+1

    def max_loop_time_decide_by_depth(self,datalen,seed_len,current_depth_number):
        """
        Decide current seed max loop time by current seed length,current datalen
        and current seed 'depth' number
        """
        if datalen >=current_depth_number: return None
        allowed_length=current_depth_number-datalen
        return max(math.floor(allowed_length/seed_len),1)

    def min_loop_time_decide_by_offset(self,datalen,seed_len,current_offset_number):
        """
        Decide current seed max loop time by current seed length,current datalen
        and current seed 'offset' number
        """
        if datalen <=current_offset_number: return False    #need 补充字节到指定长度以后再进行构造
        allowed_length=datalen-current_offset_number   #不做限制了 随意匹配
        return True

    def loop_seed(self,seed,loop_time):
        """

        :param seed: current content
        :param loop_time:
        :return: current data
        """
        data=""
        t=0
        while t<loop_time:
            data=data+seed
            t=t+1
        return data

# with open('./expRules.rules') as f:
#     rule=f.read()
# print(rule)
# #print(line)
# a=SnortRule(rule)
# print(a.opt_keyword_list)
# print(a.rule_options)
# rule_attr=SnortRuleAttr(a)
#
# print(rule_attr.option_index('content'))
# print(rule_attr.option_index('pcre'))
#
# print(rule_attr.get_opt_content())
# print(len(rule_attr.get_opt_content()))
# print(type(rule_attr.get_opt_content()[0]))
#
# print(rule_attr.get_opt_pcre())
# print(type(rule_attr.get_opt_pcre()))
# print(len(rule_attr.get_opt_pcre()))
#
# # content_and_pcre=rule_attr.get_opt_content()
# # pcre_index=rule_attr.option_index('pcre')[0]
# # index=-1
# # add_flag=False
# # for num in rule_attr.option_index('content'):
# #     index=index+1
# #     if num<pcre_index:
# #         continue
# #     else:
# #         content_and_pcre.insert(index,rule_attr.get_opt_pcre())
# #         add_flag=True
# # if not add_flag:
# #     content_and_pcre.append(rule_attr.get_opt_pcre())
# # for case in content_and_pcre:
# #     print(case)
# # construct the hard-match string
#
#
with open('./expRules.rules') as f:
    rule=f.read()
#print(rule)


c=SnortRule(rule)
d=ConstructHardMatchData(c)
#d.converge_content_and_pcre()
d.show_content_and_pcre()
d.construct_data()
d.data_add_last_pcre()
print(d.data)
print(d.datalen)

